{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ede8c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import os\n",
    "import os.path as osp\n",
    "import imageio.v2 as imageio\n",
    "\n",
    "from scipy.stats import describe\n",
    "import cv2\n",
    "# find -type f -name \"???[0-9]?.png\" ! -name \"???[0,2,4,6,8]0.png\" -exec rm -f {} +"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b3a8c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import cv2\n",
    "def imread_cv2(path, options=cv2.IMREAD_COLOR):\n",
    "    \"\"\"Open an image or a depthmap with opencv-python.\"\"\"\n",
    "    if path.endswith((\".exr\", \"EXR\")):\n",
    "        options = cv2.IMREAD_ANYDEPTH\n",
    "    img = cv2.imread(path, options)\n",
    "    if img is None:\n",
    "        raise IOError(f\"Could not load image={path} with {options=}\")\n",
    "    if img.ndim == 3:\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf8e531",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = 0\n",
    "# List of file paths\n",
    "scene_dir = \"/lc/data/3D/wildrgbd/processed/orange/scene_000\"\n",
    "file_paths = sorted([f for f in os.listdir(scene_dir) if f.endswith(\"_rgb.png\")])\n",
    "print(len(file_paths))\n",
    "cols = 2\n",
    "rows = 5\n",
    "interval = 2\n",
    "# Create a figure with 2 columns and 10 rows,dpi=100\n",
    "fig, axs = plt.subplots(rows, cols, figsize=(12, rows*4))\n",
    "# Iterate over the file paths and display the images\n",
    "for i, file_path in enumerate(file_paths[s:rows*cols*interval//2+s:interval]):\n",
    "    rgb_path = osp.join(scene_dir, file_path)\n",
    "    depth_path = rgb_path.replace(\"_rgb.png\", \"_depth.png\")\n",
    "    cam_path = rgb_path.replace(\"_rgb.png\", \"_cam.npz\")\n",
    "\n",
    "    img = imread_cv2(rgb_path, cv2.IMREAD_COLOR)\n",
    "    depth = imread_cv2(depth_path, cv2.IMREAD_UNCHANGED)/1000\n",
    "    \n",
    "    cam_file = np.load(cam_path)\n",
    "    desc = f'{depth.min()}, {depth.max()}, {depth.mean()}'\n",
    "    cam_desc = f'{cam_file['pose'][:3, -1]}'\n",
    "\n",
    "\n",
    "    axs[i * 2 // cols , i * 2 % cols ].imshow(img)\n",
    "    axs[i * 2 // cols , i * 2 % cols].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 2 // cols , i * 2 % cols].set_title(cam_desc)\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].imshow(depth)\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].set_title(desc)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "709073f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "s = 0\n",
    "# List of file paths\n",
    "root = \"/lc/data/3D/wildrgbd/original/flower_pot/scenes/scene_012/rgb\"\n",
    "file_paths = sorted([f for f in os.listdir(root) if f.endswith(\".png\")])\n",
    "cols = 3\n",
    "rows = 4\n",
    "interval = 2\n",
    "# Create a figure with 2 columns and 10 rows,dpi=100\n",
    "fig, axs = plt.subplots(rows, cols, figsize=(12, rows*5))\n",
    "\n",
    "# Iterate over the file paths and display the images\n",
    "for i, file_path in enumerate(file_paths[s:rows*cols*interval//3+s:interval]):\n",
    "    img = mpimg.imread(osp.join(root, file_path))\n",
    "    depth = cv2.imread(osp.join(root.replace(\"/rgb\",\"/depth\"), file_path), cv2.IMREAD_ANYDEPTH)\n",
    "    # depth = mpimg.imread(osp.join(root.replace(\"/rgb\",\"/depth\"), file_path))\n",
    "    mask = cv2.imread(osp.join(root.replace(\"/rgb\",\"/masks\"), file_path), cv2.IMREAD_ANYDEPTH)\n",
    "    # depth[mask==0] = 0\n",
    "    # depth[depth==0] = 1000\n",
    "    # depth = 1.0/ depth\n",
    "    axs[i * 3 // cols , i * 3 % cols ].imshow(img)\n",
    "    axs[i * 3 // cols , i * 3 % cols].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 3 // cols , i * 3 % cols].set_title(file_path.split('.')[0]+f\"_{i}\")\n",
    "    axs[i * 3 // cols , i * 3 % cols + 1].imshow(depth)\n",
    "    axs[i * 3 // cols , i * 3 % cols + 1].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 3 // cols , i * 3 % cols + 1].set_title(file_path.split('.')[0]+\"_depth\")\n",
    "    axs[i * 3 // cols , i * 3 % cols + 2].imshow(mask)\n",
    "    axs[i * 3 // cols , i * 3 % cols + 2].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 3 // cols , i * 3 % cols + 2].set_title(file_path.split('.')[0]+\"_mask\")\n",
    "\n",
    "plt.show()\n",
    "# continious frames of vedio, 20 frames about 20 degree\n",
    "# masked object, whole image has depth\n",
    "# We store depths in the depth scale of 1000. That is, when we load depth image and divide by 1000, we could get depth in meters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da44f019",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "import importlib\n",
    "import utils.rotate\n",
    "importlib.reload(utils.rotate)\n",
    "from utils.rotate import rotate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf9934a",
   "metadata": {},
   "outputs": [],
   "source": [
    "intrinsics = np.array([[436.11008,   0.     , 240.91835],\n",
    "                       [  0.     , 436.11008, 322.66653],\n",
    "                       [  0.     ,   0.     ,   1.     ]], dtype=np.float32)\n",
    "pose = np.array([[ 9.8153263e-01, -7.7415816e-02,  1.7492986e-01, -6.1428569e-02],\n",
    "                 [ 7.7405863e-02,  9.9697584e-01,  6.8903100e-03, -4.0515000e-04],\n",
    "                 [-1.7493425e-01,  6.7775301e-03,  9.8455679e-01,  1.0289070e-02],\n",
    "                 [ 0.         ,  0.         ,  0.         ,  1.         ]], dtype=np.float32)\n",
    "# 3D point in world coordinates (homogeneous)\n",
    "point_3d = np.array([0.3, 0.1, 0.5, 1.0], dtype=np.float32)\n",
    "H, W, _ = img.shape\n",
    "# Rotate data\n",
    "image_rot, depth_rot, intrinsics_new, new_pose, mask_rot = rotate(img, depth, intrinsics, pose, mask, k=1)\n",
    "\n",
    "# Step 1: Original 3D to 2D projection\n",
    "cam_point = pose @ point_3d  # 3D point in original camera coordinates\n",
    "pixel_orig = intrinsics @ cam_point[:3]  # Project to 2D\n",
    "pixel_orig = pixel_orig[:2] / cam_point[2]  # Normalize (u_orig, v_orig)\n",
    "\n",
    "# Step 2: Rotate 2D coordinates (counterclockwise, k=1)\n",
    "u_rot, v_rot = pixel_orig[1], H - 1 - pixel_orig[0]  # (u_orig, v_orig) -> (v_orig, H-1-u_orig)\n",
    "\n",
    "# Step 3: Back-project to 3D using new intrinsics and new pose\n",
    "# Use the original depth (z_c) for back-projection\n",
    "z_c = cam_point[2]  # Depth in original camera coordinates\n",
    "# 2D point in homogeneous coordinates: [u_rot, v_rot, 1] * z_c\n",
    "pixel_rot_hom = np.array([u_rot, v_rot, 1.0], dtype=np.float32) * z_c\n",
    "# Back-project to 3D camera coordinates: (x_c', y_c', z_c') = inv(intrinsics_new) @ pixel_rot_hom\n",
    "intrinsics_new_inv = np.linalg.inv(intrinsics_new)\n",
    "cam_point_new = intrinsics_new_inv @ pixel_rot_hom  # [x_c', y_c', z_c']\n",
    "# Transform to world coordinates: P_w' = inv(new_pose) @ [x_c', y_c', z_c', 1]\n",
    "cam_point_new_hom = np.append(cam_point_new, 1.0)  # Homogeneous\n",
    "new_pose_inv = np.linalg.inv(new_pose)\n",
    "world_point_new = new_pose_inv @ cam_point_new_hom\n",
    "# Transform back to original camera coordinates: P_c_back = pose @ P_w'\n",
    "cam_point_back = pose @ world_point_new\n",
    "\n",
    "\n",
    "# Compare 3D camera coordinates\n",
    "print(\"Original 3D camera coordinates (x, y, z):\", cam_point[:3])\n",
    "print(\"Back-projected 3D camera coordinates (x, y, z):\", cam_point_back[:3])\n",
    "print(\"3D Difference (should be close to zero):\", cam_point[:3] - cam_point_back[:3])\n",
    "\n",
    "# Optional: Print 2D coordinates for reference\n",
    "print(\"\\nOriginal 2D pixel coordinates (u_orig, v_orig):\", pixel_orig)\n",
    "print(\"Rotated 2D pixel coordinates (u_rot, v_rot):\", [u_rot, v_rot])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f5c5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print(depth.shape, depth.dtype, type(depth)) \n",
    "print(describe(depth, axis=None))\n",
    "print((depth==0.0).sum(), (depth>0).sum())\n",
    "\n",
    "print(depth.shape, depth.dtype, type(depth)) \n",
    "print(describe(depth[mask>0], axis=None))\n",
    "print((depth[mask>0]==0.0).sum(), (depth[mask>0]>0).sum())\n",
    "\n",
    "print(mask.shape, mask.dtype, type(mask)) \n",
    "print(describe(mask, axis=None))\n",
    "print((mask==0.0).sum(), (mask==1.0).sum())\n",
    "print(640* 480 - (mask==0.0).sum())\n",
    "\n",
    "# (640, 480) uint16 <class 'numpy.ndarray'>\n",
    "# DescribeResult(nobs=307200, minmax=(np.uint16(0), np.uint16(4995)), mean=np.float64(1689.0027766927083), variance=np.float64(1449727.8767529565), skewness=np.float64(0.28871121704302627), kurtosis=np.float64(-1.330760893130484))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2616181c",
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_to_world_path = osp.join(root.removesuffix('/rgb'), \"cam_poses.txt\")\n",
    "camera_to_world_content = np.genfromtxt(camera_to_world_path)\n",
    "camera_to_world = camera_to_world_content[:, 1:].reshape(-1, 4, 4)\n",
    "\n",
    "frame_idx = camera_to_world_content[:, 0]\n",
    "num_frames = frame_idx.shape[0]\n",
    "print(num_frames, camera_to_world.shape, camera_to_world_content.shape, type(camera_to_world_content))\n",
    "print(camera_to_world[10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45a1b34d",
   "metadata": {},
   "outputs": [],
   "source": [
    "depth = mpimg.imread(osp.join(root, file_path.replace(\"color.png\",\"depth.png\")))\n",
    "print(depth.shape, depth.dtype) \n",
    "print(describe(depth, axis=None))\n",
    "print(describe(depth*65535, axis=None))\n",
    "# import cv2\n",
    "# cv_depth = cv2.imread(osp.join(root, file_path.replace(\"color.png\",\"depth.png\")), cv2.IMREAD_ANYDEPTH)\n",
    "# print(describe(cv_depth, axis=None))\n",
    "# print(describe(cv_depth/65535, axis=None))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vggt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
